import pymysql
import json
from tabulate import tabulate
import requests
import re
import openai

openai.api_key = "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"
client = openai.OpenAI(
    api_key="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",
)

# 数据库连接信息
connection_config = {
    'host': '10.70.223.112',
    'user': 'dbuser',
    'password': 'Test_12345',
    'database': 'BI_intelligence'
}


def get_database_info():
    """获取数据库表结构和映射信息"""
    try:
        # 建立数据库连接
        connection = pymysql.connect(**connection_config)
        cursor = connection.cursor(pymysql.cursors.DictCursor)  # 使用字典游标

        # 获取字段映射信息
        cursor.execute("SELECT * FROM tdr_code_table")
        mapping_data = cursor.fetchall()

        # 获取tdr_data_table的结构信息
        cursor.execute("DESCRIBE tdr_data_table")
        table_structure = cursor.fetchall()

        # 获取tdr_data_table的样本数据
        cursor.execute("SELECT * FROM tdr_data_table LIMIT 5")
        sample_data = cursor.fetchall()

        return {
            "mapping_data": mapping_data,
            "table_structure": table_structure,
            "sample_data": sample_data
        }

    except Exception as e:
        print(f"数据库错误: {e}")
        return None

    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'connection' in locals():
            connection.close()


def create_llm_prompt(user_query, data_info):
    """创建发送给大模型的prompt，包含用户查询和数据库信息"""
    if not data_info:
        return "无法获取数据库信息"

    mapping_data = data_info["mapping_data"]
    table_structure = data_info["table_structure"]
    sample_data = data_info["sample_data"]

    # 创建字段映射字典，方便引用
    field_mappings = {}
    for item in mapping_data:
        # 假设tdr_code_table中有code和name字段
        if 'code' in item and 'name' in item:
            field_mappings[item['code']] = item['name']

    # 格式化字段映射信息
    mapping_info = "字段映射信息 (字段名 -> 中文含义):\n"
    for code, name in field_mappings.items():
        mapping_info += f"- {code} -> {name}\n"

    # 格式化表结构信息
    structure_info = "tdr_data_table表结构:\n"
    for field in table_structure:
        field_name = field.get('Field', '')
        chinese_name = field_mappings.get(field_name, '无映射')
        structure_info += f"- 字段名: {field_name}, 中文含义: {chinese_name}, 类型: {field.get('Type', '')}, 可空: {field.get('Null', '')}\n"

    # 格式化样本数据
    sample_info = "数据样本:\n"
    if sample_data:
        # 使用tabulate生成表格样式的样本数据
        headers = sample_data[0].keys()
        rows = [[item[col] for col in headers] for item in sample_data]
        sample_table = tabulate(rows, headers=headers, tablefmt="grid")
        sample_info += sample_table

    # 组合prompt
    prompt = f"""
我有一个名为tdr_data_table的数据表，下面是该表的相关信息：

{structure_info}

{mapping_info}

{sample_info}

用户的自然语言查询是: "{user_query}"

请将用户的自然语言查询转换为针对tdr_data_table表的SQL查询语句。生成的SQL应该正确反映用户的查询意图。

请只返回SQL语句，不需要解释。
"""
    return prompt


def get_sql_from_llm(prompt, api_key=None):
    """向大模型发送请求并获取SQL查询语句"""
    print("\n准备发送给大模型的Prompt:\n")
    print(prompt)
    print("\n====================================================\n")

    # 这里是使用大模型API的示例代码，需要根据实际情况修改
    # 下面提供了几种常见大模型API的调用方式，取消注释你要使用的部分，并填入API密钥

    response = client.chat.completions.create(
        model="gpt-4o",  # 或 "gpt-3.5-turbo"
        messages=[
            {"role": "system",
             "content": "你是一个专业的SQL查询生成器。请根据提供的数据库信息，将用户的自然语言查询转换为SQL查询语句。"},
            {"role": "user", "content": prompt}
        ],
        temperature=0.1  # 较低的温度使输出更加确定性
    )
    res = response.choices[0].message.content.strip()
    cleaned_query = res.replace("```sql\n", "")
    cleaned_query = cleaned_query.replace(";\n```", "")
    return cleaned_query


def execute_sql(sql_query):
    """执行SQL查询并返回结果"""
    try:
        connection = pymysql.connect(**connection_config)
        cursor = connection.cursor(pymysql.cursors.DictCursor)

        print(f"执行SQL: {sql_query}")
        cursor.execute(sql_query)
        results = cursor.fetchall()

        if results:
            headers = results[0].keys()
            rows = [[item[col] for col in headers] for item in results]
            result_table = tabulate(rows, headers=headers, tablefmt="grid")
            return f"查询结果:\n{result_table}"
        else:
            return "查询没有返回任何结果"

    except Exception as e:
        return f"SQL执行错误: {e}"

    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'connection' in locals():
            connection.close()


def main():
    print("正在获取数据库信息...")
    data_info = get_database_info()

    if not data_info:
        print("获取数据库信息失败，请检查连接配置")
        return

    print("\n数据库信息获取成功！")
    print(f"映射表中有 {len(data_info['mapping_data'])} 条记录")
    print(f"数据表有 {len(data_info['table_structure'])} 个字段")

    while True:
        print("\n" + "=" * 60)
        user_query = input("请输入你的自然语言查询 (输入'exit'退出): ")

        if user_query.lower() == 'exit':
            print("程序已退出")
            break

        # 创建prompt并获取SQL
        prompt = create_llm_prompt(user_query, data_info)
        sql_query = get_sql_from_llm(prompt)

        print("\n生成的SQL查询:")
        print(sql_query)

        # 询问是否执行查询
        execute = input("\n是否执行此SQL查询? (y/n): ").lower() == 'y'
        if execute:
            result = execute_sql(sql_query)
            print(result)


if __name__ == "__main__":
    main()